Sorry, I missed this one

Merged — thanks, this was a substantial piece of work and it holds up well under review.

The seven fix(dspark)/test/chore commits went in as 1a2e039039, and ship on
codex/ds4-sm120-min-enable + ds4-sm120-preview-dev as part of 1ee1ec7930
(which also carries 34 upstream commits merged on top).

Validation on GB10 (2× SM121a, TP=2, MTP2, mml 49152, util 0.85, fp8 KV)

Unit suite green: 795 kv_offload, 6 DSv4 sparse-SWA/ubatch/indexer, 35 rejection sampler.
A DSpark serve answers correctly (17×23 → 391, capital of France, first three primes).

Re-gated after merging 34 upstream commits on top: arthur c=1 2/2, IMA 0, and GSM8K three times
— 0.9469 / 0.9507 / 0.9583. Worth stating plainly since it bears on the numbers above: that
1.1pp spread on one unchanged build is wider than our entire historical build-to-build range, so
GSM8K here discriminates gross correctness breakage, not sub-1pp quality deltas.

Kernel-test A/B at 520c25949c vs the merge, same node and environment:
test_mhc_kernels.py 13 → 11 failures, test_rejection_sampler_utils.py 29 → 31 passed
(your two new tests pass). The one AFTER-only failure,
test_mhc_fused_post_pre_reuses_dead_buffers[128], is the [128] parametrization that already
fails for the pre-existing test_mhc_fused_post_pre[*-128] without the reuse variant — it
inherits an existing failure rather than introducing one. No new failure modes.

Review notes

The parts that needed the most scrutiny were the in-place mutations, since geometry- and
aliasing-invariant bugs have bitten this branch repeatedly. All of them check out:

• Rewriting the C128A prefill indices in place local→global and caching them across layers
  is safe: the Triton kernel reads and writes the same element of each row with no cross-row
  reads, and the super()._forward_prefill fallback — whose parent reads that field expecting
  local indices — is gated on the global VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL, so
  the whole model takes one path and layers never mix readers.
• Sharing one c128a_topk_buffer between decode and prefill partitions exactly on
  is_decode = token_idx < num_decode_tokens; the stale tail of the shared lens buffer is
  harmless because the consumer stores topk_lens[token_idx] unconditionally.
• Making _mtp_hidden_buffer conditional is fine for DeepSeek-V4-Flash: its
  inference/config.json has both compress_ratios and hc_mult (4), so the MTP path still
  allocates it.

Two deltas relative to your current branch head

1. rejection_sampler.py — we bind the trailing optionals by keyword. When we merged, the
draft_logits_index_mapping threading was not yet present, so self.synthetic_conditional_rates
was passed positionally into the slot rejection_sample() reserves for
draft_logits_index_mapping — the two are adjacent in the signature, so it bound silently with
no conflict and no error. You have since fixed the same thing by adding the missing argument;
ours is the same semantics with keyword binding plus a comment, which is more robust if that
signature grows again. Functionally equivalent — just flagging that we kept ours.

2. fix(config): keep sequence-parallel MoE enabled at DP=1 — we did not take this one.
Not a disagreement with your measurements, but it doesn't fit our branch:

• It's inert for us. use_sequence_parallel_moe also requires enable_expert_parallel, and
  expert parallel is off in our standard GB10 serve, so the property is False either way.
• Reverting an upstream commit means carrying a conflict point through every future upstream
  merge, and silently re-breaking the pure-TP memory case Fix: Restore data_parallel_size > 1 for use_sequence_parallel_moe vllm-project/vllm#48849 was written for, for everyone
  using this fork.
• If we later enable EP at DP=1, flipping use_sequence_parallel_moe to True also flips
  pad_shared_expert to False in both dspark.py and model.py — the DSpark shared-expert
  weight-loading path that broke a 4-node deployment a week ago.

Your own commit message points at the right answer: make the gate config-aware rather than
requiring DP > 1. That's an upstream conversation, and your numbers are good ammunition for it —
worth taking to vllm-project#48849 directly.

Two things worth knowing

• mhc_fused_post_pre_tilelang_reuse_residual is not direct_register_custom_op'd and is
  absent from __all__, unlike the mhc_fused_post_pre_tilelang it replaces on the hot path.
  We measured no effect (0 graph breaks, CUDA-graph capture 24 s), and it is arguably the safer
  side — an unregistered plain function leaves its in-place mutation visible to the tracer.
  Please don't "fix" this by registering it under the file's uniform mutates_args=[]:
  that would declare a mutation that does happen as absent.
• We could not confirm the VRAM saving on our hardware, and we're not disputing it. The
  serve-reported "available KV memory / GPU KV cache size" turned out to swing ~1 GiB / 9%
  between runs of near-identical builds on our setup, which is larger than the effect we were
  trying to resolve, so single-run comparisons of that figure are not meaningful here. Our
  default path is MTP, which keeps _mtp_hidden_buffer and so doesn't collect the DSpark win
  anyway.

Closing as merged. Thanks again.